//
// Created by hujin on 2020/11/8.
//

#include <cassert>
#include <tuple>
#include <string.h>
#include "mem_allocator.h"
#include "../logger.h"

std::string getDataName(int varid){ // global name
    std::string ret = "var_";
    ret.append(std::to_string(varid));
    return ret;
}
int MemoryAllocator::requestReg(int desire, bool global) {
    int reg = global ? registerState->requestGlobalReg(desire) : registerState->requestNewReg(desire);
    if(reg >= 0) return reg;
    else {
        reg = registerState->releaseAndReuseReg();
        assert(reg > 0);
        if(regCur[reg] != 0){
            save(regCur[reg]);
            data[regCur[reg]].isReg = false;
            //releaseReg(regCur[reg]);
        }
        return reg;
    }
}



void MemoryAllocator::allocGlobalMem(int varid, int space) {
    data.insert({varid, {false, true, -1, 0, false}});
    synthesizer->data<<getDataName(varid)<<":.space "<<space<<std::endl;
}
void MemoryAllocator::allocGlobalMem(int varid, std::vector<int> init) {
    data.insert({varid, {false, true, -1, 0, false}});
    synthesizer->data<<getDataName(varid)<<":.word ";
    for(int i = 0; i < init.size(); i++){
        if(i > 0)
            synthesizer->data<<',';
        synthesizer->data<<init[i];
    }
    synthesizer->data<<std::endl;
}


void MemoryAllocator::allocLocalArr(int varid, int space) {
    int ad = stackMemAlloc(space);
    data.insert({varid, {false, false, -1, ad, false}});

}

int MemoryAllocator::stackMemAlloc(int size) {
    stack_length = stack_length + size;
    return  stack_length - size;
}
int MemoryAllocator::loadToReg_writeOnly(int varid_old){
    return _loadToReg(varid_old, true, true);
}
int MemoryAllocator::loadToReg(int varid_old, bool write /*TODO: check read*/) {
    return _loadToReg(varid_old, write, false);
}

int MemoryAllocator::_loadToReg(int varid_old, bool write, bool notRead) {
    int varid = redirect.count(varid_old) ? redirect[varid_old] : varid_old;//find_ancestor(redirect, varid_old);
    if(write && varid != varid_old) {
        synthesizer->text << "# kill redirect of var" << varid_old << " to var" << varid << std::endl;
        redirect.erase(varid_old);
        int regT= requestReg(data[varid_old].regId, varid_old > 0);
        auto &ad  = data[varid_old];
        ad.isReg = true;
        ad.regId = regT;
        synthesizer->text << "# link reg$" << regT << " to var" << varid_old << std::endl;
        regCur[regT] = varid_old;
        ad.valChanged = true;
        //TODO:extract
        if(!notRead) {
            if (data[varid].isReg) {
                synthesizer->text << "move $" << regT << " $" << data[varid].regId << std::endl;
            } else {
                if (data[varid].isGlobal) synthesizer->text << "lw $" << regT << ' ' << getDataName(varid) << std::endl;
                else synthesizer->text << "lw $" << regT << ' ' << -(data[varid].dataAddr) << "($sp)" << std::endl;
            }
        }

        return regT;
    }
    if(!data.count(varid)){
        int reg = requestReg(0, varid_old > 0);
        data.insert({varid, {true, false, reg, stackMemAlloc(4), write}});
        regCur[reg] = varid;
        synthesizer->text << "# link reg$" << reg << " to var" << varid << std::endl;
        return reg;
    }
    auto &ad = data[varid];
    ad.valChanged |= write;
    if(ad.isReg) {
        registerState->markUsed(ad.regId);
        return ad.regId;
    }
    ad.isReg = true;
    int reg = requestReg(ad.regId, varid_old > 0);
    ad.regId = reg;
    regCur[reg] = varid;
    synthesizer->text << "# link reg$" << reg << " to var" << varid << std::endl;
    if(!notRead) {
        if (ad.isGlobal) synthesizer->text << "lw $" << reg << ' ' << getDataName(varid) << std::endl;
        else synthesizer->text << "lw $" << reg << ' ' << (-ad.dataAddr) << "($sp)" << std::endl;
    }
    return reg;
}

void MemoryAllocator::move(int varids, int varidt) {
    int realS = redirect.count(varids) ? redirect[varids] : varids;
    if(realS == varidt)return;
    if(varidt < 0){ // will not merge
        if(!data.count(varidt)){
            data.insert({varidt, {false, false, -1, stackMemAlloc(4), false}});
        }
        releaseReg(varidt);
        synthesizer->text<<"# redirect var"<<varidt<<" to var"<<realS<<std::endl;
        redirect[varidt] = realS;
    } else{
        int regT = loadToReg_writeOnly(varidt); // TODO 优化:直接存进内存？
        int regS = loadToReg(realS, false);
        synthesizer->text<<"move $"<<regT<<" $"<<regS<<std::endl;
    }
}


//oldt: free
//read t -> read s
//since reg(t) is free, saving will not occur in releaseReg
//write t -> kill link, write to t;

//void MemoryAllocator::move(int varids, int varidt) {
//    int realS = redirect.count(varids) ? redirect[varids] : varids;
//    if(realS == varidt)return;
//    releaseReg(varidt);
//    redirect[varidt] = realS;
//}

void MemoryAllocator::releaseReg(int varid) {
    if(data.count(varid) && data[varid].isReg){
        synthesizer->text<<"#release reg$"<<data[varid].regId<<" by var"<<varid<<std::endl;
        regCur[data[varid].regId] = 0;
        data[varid].isReg = false;
        registerState->releaseReg(data[varid].regId);
    }
}
void MemoryAllocator::useReg(int regNum){
    if(regCur[regNum] > 0){
        save(regCur[regNum]);
        releaseReg(regCur[regNum]);
    }

}
void MemoryAllocator::save(int varid) {
    int source = redirect.count(varid) ? redirect[varid] : varid;
    if(!data.count(varid)){
        log_error(5, "mem_allocator.h", 77, 0);//save before create
        return;
    }
    auto &ad = data[varid];
    if(source != varid) {
        int reg = loadToReg(source, false);
        if(ad.isGlobal) synthesizer->text<<"sw $"<<reg<<' ' <<getDataName(varid)<<std::endl;
        else synthesizer->text<<"sw $"<<reg<<' '<<(-ad.dataAddr)<<"($sp)"<<std::endl;

        return ;
    }
    if(ad.isReg && ad.valChanged) {
        if(ad.isGlobal) synthesizer->text<<"sw $"<<ad.regId<<' ' <<getDataName(varid)<<std::endl;
        else synthesizer->text<<"sw $"<<ad.regId<<' '<<-ad.dataAddr<<"($sp)"<<std::endl;
        ad.valChanged = false;
    }
}

int MemoryAllocator::loadArrToReg(int varidarr, int varidtarget, int offset) {
    int regT = loadToReg_writeOnly(varidtarget);
    auto d = data[varidarr];
    if(d.isGlobal)synthesizer->text<<"lw $"<<regT<<' '<<getDataName(varidarr)<<'+'<<offset<<std::endl;
    else synthesizer->text<<"lw $"<<regT<<' '<<-(offset + d.dataAddr)<<"($sp)"<<std::endl;
    return regT;
}

void MemoryAllocator::saveArr(int varidarr, int varid_source, int offset) {
    int regT = loadToReg(varid_source, false);
    auto d = data[varidarr];
    if(d.isGlobal)synthesizer->text<<"sw $"<<regT<<' '<<getDataName(varidarr)<<'+'<<offset<<std::endl;
    else synthesizer->text<<"sw $"<<regT<<' '<<-(offset + d.dataAddr)<<"($sp)"<<std::endl;
}




int MemoryAllocator::loadArrToRegVarPos(int varidarr, int varidTarget, int offsetVarid) {
    int regT = loadToReg_writeOnly(varidTarget);
    auto d = data[varidarr];
    int regI = loadToReg(offsetVarid, !d.isGlobal);
    if(d.isGlobal)synthesizer->text<<"lw $"<<regT<<' '<<getDataName(varidarr)<<"($"<<regI<<")"<<std::endl;
    else{
        synthesizer->text<<"subu $"<<regI<<" $sp"<<" $"<<regI<<std::endl;
        synthesizer->text<<"lw $"<<regT<<' '<<-d.dataAddr<<"($"<<regI<<")"<<std::endl;
    }
    return regT;
}

void MemoryAllocator::saveArrVarPos(int varidarr, int varid_source, int offsetVarid) {
    int regT = loadToReg(varid_source, false);
    auto d = data[varidarr];
    int regI = loadToReg(offsetVarid, !d.isGlobal);
    if(d.isGlobal)synthesizer->text<<"sw $"<<regT<<' '<<getDataName(varidarr)<<"($"<<regI<<")"<<std::endl;
    else{
        synthesizer->text<<"subu $"<<regI<<" $sp"<<" $"<<regI<<std::endl;
        synthesizer->text<<"sw $"<<regT<<' '<<-d.dataAddr<<"($"<<regI<<")"<<std::endl;
    }
}

MemoryAllocator *MemoryAllocator::split() {
    auto *allocator = new MemoryAllocator(synthesizer, nullptr);
    allocator->stack_length = stack_length;
    allocator->data = this->data;
    allocator->redirect = redirect;
    allocator->registerState = registerState->split();
    for(int reg = 0; reg < reg_cnt; reg++) allocator->regCur[reg] = regCur[reg];
    return allocator;
}

void MemoryAllocator::mergeTo(MemoryAllocator *memNext, const std::set<int>& activeIds) {
    int regTemp = reg_temp_0;

    synthesizer->text<<"# merge start"<<std::endl;
    for(int reg = 0; reg < reg_cnt; reg++){
        int id = memNext->regCur[reg];
        if(id > 0 && /* TODO activeIds.count(id) &&*/ regCur[reg] != id){
            bool occupied = memNext->data.count(regCur[reg]) && memNext->data[regCur[reg]].isReg;
            if(occupied /* TODO && activeIds.count(regCur[reg])*/ && data[id].isReg){
                //swap
                int anotherId = regCur[reg];
                synthesizer->text << "move $" << regTemp << " $" << reg<<std::endl;
                regCur[regTemp] = anotherId;
                regCur[data[anotherId].regId] = 0;

                std::swap(data[anotherId].regId ,  regTemp); // like hanoi, regTemp -> transfer tower;

                moveToSpecialReg(id, reg);
                //synthesizer->text << "move $" << reg << " $" << data[id].regId<<std::endl;
                data[id].regId = reg;
                regCur[reg] = data[id].regId;
            } else if(occupied){
                int anotherId = regCur[reg];
                moveToSpecialReg(anotherId, memNext->data[anotherId].regId);
                data[anotherId].regId = memNext->data[anotherId].regId;
                regCur[data[anotherId].regId] = anotherId;

                moveToSpecialReg(id, reg);
                data[id].isReg = true;
                data[id].regId = reg;
                regCur[reg] = data[id].regId;
            }
            else {
                if(regCur[reg] > 0){
                    save(regCur[reg]);
                    releaseReg(reg);
                }
                moveToSpecialReg(id, reg);
                data[id].isReg = true;
                data[id].regId = reg;
                regCur[reg] = data[id].regId;
            }
        }
        int idC = regCur[reg];
        if(idC > 0 && memNext->data.count(idC) &&  !memNext->data[idC].isReg){
            save(idC);
        }
    }
    synthesizer->text<<"# merge end"<<std::endl;
    regCur[reg_temp_0] = 0;
}

void MemoryAllocator::saveArrCnst(int varidarr, int value, int offset) {
    int regT = reg_temp_0;
    synthesizer->text<<"li $"<<regT<<' '<<value<<std::endl;
    auto d = data[varidarr];
    if(d.isGlobal)synthesizer->text<<"sw $"<<regT<<' '<<getDataName(varidarr)<<'+'<<offset<<std::endl;
    else synthesizer->text<<"sw $"<<regT<<' '<<-(offset + d.dataAddr)<<"($sp)"<<std::endl;
}

void MemoryAllocator::saveArrCnstVarPos(int varidarr, int value, int offsetVarid) {
    int regT = reg_temp_0;
    synthesizer->text<<"li $"<<regT<<' '<<value<<std::endl;
    auto d = data[varidarr];
    int regI = loadToReg(offsetVarid, !d.isGlobal);
    if(d.isGlobal)synthesizer->text<<"sw $"<<regT<<' '<<getDataName(varidarr)<<"($"<<regI<<")"<<std::endl;
    else{
        synthesizer->text<<"subu $"<<regI<<" $sp"<<" $"<<regI<<std::endl;
        synthesizer->text<<"sw $"<<regT<<' '<<-d.dataAddr<<"($"<<regI<<")"<<std::endl;
    }
}

void MemoryAllocator::setToSpecialReg(int varid, int regNum) {
    if(regCur[regNum] > 0) {
        save(regCur[regNum]);
        releaseReg(regCur[regNum]);
    }
    if(!data.count(varid)){
        data.insert({varid, {true, false, regNum, stackMemAlloc(4), false}});
        regCur[regNum] = varid;
    } else{
        data[varid].isReg = true;
        data[varid].regId = regNum;
        regCur[regNum] = varid;
    }

}

void MemoryAllocator::moveToSpecialReg(int varid_old, int regNum) {
    int varid = redirect.count(varid_old) ? redirect[varid_old] : varid_old;//find_ancestor(redirect, varid_old);
    if(!data.count(varid)){
        log_error(5, "mem_allocator.h", 222, 0);//save before create
    }
    auto ad = data[varid];
    if(ad.isReg){
        if(ad.regId != regNum) synthesizer->text<<"move $"<<regNum<<" $"<<ad.regId<<std::endl;
    } else{
        if(ad.isGlobal) synthesizer->text<<"lw $"<<regNum<<' ' <<getDataName(varid)<<std::endl;
        else synthesizer->text<<"lw $"<<regNum<<' '<<-ad.dataAddr<<"($sp)"<<std::endl;
    }
}

MemoryAllocator::~MemoryAllocator() {
    delete registerState;
}
void MemoryAllocator::saveRegs(const std::set<int>& activeIds){
    for (int i : regCur) {
        if(i != 0 /*TODO && activeIds.count(i)*/){
            save(i);
            //releaseReg(i);
        }
    }
}
void MemoryAllocator::push_stack() {
    if(param_cnt == 0) stack_length_freeze = stack_length;//not frozen
    synthesizer->text<<"sw $31, "<<-stack_length_freeze<<"($sp)"<<std::endl;
    stack_length_freeze += 4;
    for(int i = 0; i < reg_cnt; i++) {
        if( regCur[i] != 0 /*TODO data flow*/) {
            releaseReg(regCur[i]);
        }
    }
    param_cnt = 0;
    if(stack_length_freeze != 0)synthesizer->text<<"addi $sp $sp, "<<-stack_length_freeze<<std::endl;
}

void MemoryAllocator::pop_stack() {
    stack_length = stack_length_freeze;
    if(stack_length != 0)synthesizer->text<<"addi $sp $sp, "<<stack_length<<std::endl;
    stack_length -= 4;
    synthesizer->text<<"lw $31, "<<-stack_length<<"($sp)"<<std::endl;
}

void MemoryAllocator::loadToParam(int varid, int maxParamCnt) {
    //assert(varid > 0);
    if(param_cnt == 0) {
        //start loading calling process
        stack_length_freeze = stack_length;
        if(maxParamCnt > reg_func_call_cnt){
            stack_length += (maxParamCnt - reg_func_call_cnt) * 4;
        }
    }
    if(param_cnt < reg_func_call_cnt){
        moveToSpecialReg(varid, getRegax(param_cnt));
        param_cnt++;
    } else {
        //data.insert({varid, {true, false, reg, stack_length_freeze + (param_cnt - reg_func_call_cnt)*4, true}});
        //regCur[reg] = varid;
        data[varid].dataAddr = (param_cnt - reg_func_call_cnt)*4;
        save(varid);
        releaseReg(varid);
        //synthesizer->text<<"# link reg$"<<reg<<" to var"<<varid<<std::endl;
        param_cnt++;
    }
}

void MemoryAllocator::readFromParam(int varid, int maxParamCnt) {
    if(read_param_cnt == 0 && maxParamCnt > reg_func_call_cnt){
        stack_length += (maxParamCnt - reg_func_call_cnt) * 4;
    }
    if(read_param_cnt < reg_func_call_cnt){
        data[varid] = {true,  false, getRegax(read_param_cnt), stackMemAlloc(4), true};
        regCur[getRegax(read_param_cnt)] = varid;
        read_param_cnt++;
    } else{
        data[varid] = {false,  false, 0, (read_param_cnt - reg_func_call_cnt) * 4, false};
        read_param_cnt++;
    }
}

void MemoryAllocator::setSynthesizer(Synthesizer *s) {
    MemoryAllocator::synthesizer = s;
}

void MemoryAllocator::save_global_changes(){
    for(auto x : data){
        if(x.second.isGlobal && x.second.isReg && x.second.valChanged){
            save(x.first);
        }
    }
}










